import numpy as np
import matplotlib.pyplot as plt
from scipy.spatial import Voronoi, voronoi_plot_2d
from sklearn.neighbors import KDTree
from shapely.geometry import Polygon, Point
import networkx as nx
from heapq import heappop, heappush
import cv2
import numpy as np
import matplotlib.pyplot as plt
import networkx as nx

def astar(graph, start, goal):
    def heuristic(node):
        return np.linalg.norm(np.array(node) - np.array(goal))

    frontier = [(0, start)]
    came_from = {start: None}
    cost_so_far = {start: 0}

    while frontier:
        current_cost, current_node = heappop(frontier)

        if current_node == goal:
            path = []
            while current_node is not None:
                path.append(current_node)
                current_node = came_from[current_node]
            return path[::-1]

        for neighbor in graph.neighbors(current_node):
            new_cost = cost_so_far[current_node] + graph[current_node][neighbor]['weight']
            if neighbor not in cost_so_far or new_cost < cost_so_far[neighbor]:
                cost_so_far[neighbor] = new_cost
                priority = new_cost + heuristic(neighbor)
                heappush(frontier, (priority, neighbor))
                came_from[neighbor] = current_node

    return None

def createGridTest():
    # 创建栅格地图
    grid_size = 250
    grid_map = np.zeros((grid_size, grid_size))

    # 设置关键点
    key_points = [(10, 10), (30, 32), (70, 78), (80, 1)]

    # 标记关键点
    for point in key_points:
        grid_map[point] = 1  # 1 表示关键点

    # 构建栅格图
    G = nx.grid_2d_graph(grid_size, grid_size)

    # 添加边的权重
    for edge in G.edges():
        G[edge[0]][edge[1]]['weight'] = 1

    # 运行A*算法连接关键点
    path = []
    for i in range(len(key_points) - 1):
        start, goal = key_points[i], key_points[i + 1]
        path += astar(G, start, goal)[:-1]

    # 绘制栅格地图和路径
    plt.imshow(grid_map, cmap='viridis', origin='lower')

    for point in key_points:
        plt.text(point[1], point[0], 'X', color='red', ha='center', va='center')

    for i in range(len(path) - 1):
        plt.plot([path[i][1], path[i + 1][1]], [path[i][0], path[i + 1][0]], color='blue', linewidth=2)

    plt.show()


class GetRoad():
    def __init__(self):
        self.file = './resource/demo2.png'
        self.fromWeb = True

    def getRoad(self):

        # 读取图片
        image = cv2.imread(self.file)

        # 将图片转换为RGB格式
        image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)


        # 设定颜色阈值，提取道路区域
        lower_bound = np.array([0, 0, 0], dtype=np.uint8)
        upper_bound = np.array([120, 120, 120], dtype=np.uint8)
        road_mask = cv2.inRange(image_rgb, lower_bound, upper_bound)



        # 进行腐蚀和膨胀操作，去除噪声
        kernel = np.ones((3, 3), np.uint8)
        road_mask = cv2.morphologyEx(road_mask, cv2.MORPH_OPEN, kernel)
        road_mask = cv2.morphologyEx(road_mask, cv2.MORPH_CLOSE, kernel)

        kernel = np.ones((5, 5), np.uint8)
        road_mask2 = cv2.dilate(road_mask, kernel, iterations=5)#膨胀
        road_mask3 = cv2.erode(road_mask2, np.ones((5, 5), np.uint8), iterations=6)#腐蚀

        road_mask = road_mask - road_mask3
        cv2.imshow("ll",road_mask3)

        # Canny边缘检测
        edges = cv2.Canny(road_mask, 50, 150)

        # 绘制原始图片和处理结果
        plt.subplot(1, 3, 1)
        plt.imshow(image_rgb)
        plt.title('Original Image')

        plt.subplot(1, 3, 2)
        plt.imshow(road_mask, cmap='gray')
        plt.title('Road Mask')

        plt.subplot(1, 3, 3)
        plt.imshow(edges, cmap='gray')
        plt.title('Edge Detection')

        plt.show()

    def getRoad2(self):
        # 读取道路图像
        road_image = cv2.imread(self.file, cv2.IMREAD_COLOR)

        # 二值化图像
        # _, binary_image = cv2.threshold(road_image, 0, 200, cv2.THRESH_BINARY)

        # 将图片转换为RGB格式

        # 设定颜色阈值，提取道路区域
        lower_bound = np.array([0, 0, 0], dtype=np.uint8)
        upper_bound = np.array([120, 120, 120], dtype=np.uint8)
        road_mask = cv2.inRange(road_image, lower_bound, upper_bound)

        # 进行腐蚀和膨胀操作，去除噪声
        kernel = np.ones((3, 3), np.uint8)
        road_mask = cv2.morphologyEx(road_mask, cv2.MORPH_OPEN, kernel)
        road_mask = cv2.morphologyEx(road_mask, cv2.MORPH_CLOSE, kernel)

        kernel = np.ones((5, 5), np.uint8)
        road_mask2 = cv2.dilate(road_mask, kernel, iterations=5)  # 膨胀
        binary_image = cv2.erode(road_mask2, np.ones((5, 5), np.uint8), iterations=6)  # 腐蚀


        cv2.imshow("ll",binary_image)

        # 寻找轮廓
        contours, _ = cv2.findContours(binary_image, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        # print(contours)
        imgCnts = np.zeros(binary_image.shape[:2], np.uint8)
        imgCnts = cv2.drawContours(imgCnts, contours, -1, (255, 255, 255), thickness=2)
        cv2.imshow("ll",imgCnts)

        # 创建空的图
        G = nx.Graph()

        # 遍历轮廓
        for contour in contours:
            # 对轮廓进行简化
            epsilon = 0.02 * cv2.arcLength(contour, True)
            approx = cv2.approxPolyDP(contour, epsilon, True)

            # 将轮廓点添加为节点
            for point in approx:
                x, y = point[0]
                G.add_node((x, y))

            # 将轮廓相邻的点之间添加边
            for i in range(len(approx) - 1):
                point1 = tuple(approx[i][0])
                point2 = tuple(approx[i + 1][0])
                G.add_edge(point1, point2)

            # 链接轮廓首尾点
            G.add_edge(tuple(approx[0][0]), tuple(approx[-1][0]))

        # 可视化图
        h,w = road_image.shape[0:2]
        plt.figure(figsize=(8, 8))
        pos = {node: (node[1], -node[0]) for node in G.nodes()}  # 调整坐标方向
        nx.draw(G, pos, node_size=10, with_labels=False, edge_color='b', node_color='r')
        plt.title('Extracted Road Network')
        plt.show()


    def getGuJia(self, file = './resource/demo2.png'):
        def skeletonize(image):
            size = np.size(image)
            skel = np.zeros(image.shape, np.uint8)

            ret, img = cv2.threshold(image, 127, 255, 0)
            element = cv2.getStructuringElement(cv2.MORPH_CROSS, (3, 3))
            done = False

            while not done:
                eroded = cv2.erode(img, element)
                temp = cv2.dilate(eroded, element)
                temp = cv2.subtract(img, temp)
                skel = cv2.bitwise_or(skel, temp)
                img = eroded.copy()

                zeros = size - cv2.countNonZero(img)
                if zeros == size:
                    done = True
            return skel

        # 读取道路图像
        road_image = cv2.imread(file, cv2.IMREAD_COLOR)
        # road_image = cv2.blur(road_image, (3, 3), 0)

        def percentage_filter(image, neighborhood_size=3):
            # 使用膨胀操作获取邻域内的最大值
            max_filtered = cv2.erode(image, np.ones((neighborhood_size, neighborhood_size), np.uint8))

            return max_filtered

        road_image2 = percentage_filter(road_image, neighborhood_size=3)

        # 二值化图像
        # 设定颜色阈值，提取道路区域
        lower_bound = np.array([0, 0, 0], dtype=np.uint8)
        upper_bound = np.array([120, 120, 120], dtype=np.uint8)
        road_mask = cv2.inRange(road_image, lower_bound, upper_bound)

        # kernel = np.ones((3, 3), np.uint8)
        # road_mask = cv2.dilate(road_mask, kernel, iterations=4)  # 膨胀
        # road_mask = cv2.ximgproc.thinning(road_mask)


        cv2.imshow("PP",road_mask)

        # 进行腐蚀和膨胀操作，去除噪声
        kernel = np.ones((3, 3), np.uint8)
        road_mask = cv2.morphologyEx(road_mask, cv2.MORPH_OPEN, kernel)
        road_mask = cv2.morphologyEx(road_mask, cv2.MORPH_CLOSE, kernel)

        kernel = np.ones((5, 5), np.uint8)
        road_mask2 = cv2.dilate(road_mask, kernel, iterations=5)  # 膨胀
        # road_mask2 = cv2.GaussianBlur(road_mask2, (7, 7), 0)

        binary_image = cv2.erode(road_mask2, np.ones((3, 3), np.uint8), iterations=7)  # 腐蚀




        # 提取图像骨架
        # skeleton = skeletonize(binary_image)
        if self.fromWeb :
            skeleton = cv2.ximgproc.thinning(road_mask)#如果图像来自前端绘图，则直接使用，无需膨胀腐蚀
        else:
            skeleton = cv2.ximgproc.thinning(binary_image)
        #cv2.imshow("PP",skeleton)


        return skeleton
    def guJia1(self):
        def get_skeleton_graph(skeleton):
            graph = nx.Graph()
            rows, cols = np.where(skeleton > 0)

            # 构建骨架图的邻接关系
            for i in range(len(rows)):
                neighbors = []

                # 寻找相邻像素
                for dr in range(-1, 2):
                    for dc in range(-1, 2):
                        if dr == 0 and dc == 0:
                            continue
                        nr, nc = rows[i] + dr, cols[i] + dc
                        if (nr, nc) in zip(rows, cols) and (nr, nc) < (rows[i], cols[i]):
                            neighbors.append((nr, nc))

                # 添加节点和边
                current_node = (cols[i], rows[i])
                graph.add_node(current_node)
                graph.add_edges_from([(current_node, neighbor) for neighbor in neighbors])

            return graph

        def split_skeleton_into_segments(skeleton_graph):
            segments = list(nx.connected_components(skeleton_graph))

            # 可以选择性地将每个连通分量绘制成一个子图
            plt.figure(figsize=(12, 6))
            for i, segment in enumerate(segments):
                plt.subplot(1, len(segments), i + 1)
                nx.draw(skeleton_graph.subgraph(segment), with_labels=False, node_size=10, edge_color='b',
                        node_color='r')
                plt.title(f'Segment {i + 1}')

            plt.show()
        # 提取图像骨架
        skeleton = self.getGuJia(self.file)

        cv2.imshow("PP",skeleton)

        cv2.waitKey(0)

        # 获取骨架图
        # skeleton_graph = get_skeleton_graph(skeleton)

        # 可视化图
        # plt.figure(figsize=(8, 8))
        # pos = {node: (node[0], -node[1]) for node in skeleton_graph.nodes()}  # 调整坐标方向
        # nx.draw(skeleton_graph, pos, node_size=10, with_labels=False, edge_color='b', node_color='r')
        # plt.title('Extracted Road Network (Skeleton)')
        # plt.show()
        # 分割骨架图成连通分量
        # split_skeleton_into_segments(skeleton_graph)

    #基于距离分割路线
    def fenShuiLin(self):
        from scipy import ndimage as ndi
        from skimage import morphology, feature
        from skimage.segmentation import watershed

        skeleton = self.getGuJia('./resource/demo2.png')
        distance = ndi.distance_transform_edt(skeleton)  # 距离变换

        cv2.imshow('pp',distance)
        local_maxi = feature.peak_local_max(distance,  footprint=np.ones((3, 3)),
                                            labels=skeleton)  # 寻找峰值
        print(local_maxi.shape)
        markers = ndi.label(local_maxi)[0]  # 初始标记点

        print(skeleton.shape)
        print(distance.shape)
        print(markers.shape)
        # labels = cv2.watershed(skeleton, markers)

        labels = watershed(-distance,  mask=skeleton)  # 基于距离变换的分水岭算法

        fig, axes = plt.subplots(nrows=2, ncols=2, figsize=(8, 8))
        axes = axes.ravel()
        ax0, ax1, ax2, ax3 = axes

        ax0.imshow(skeleton, cmap=plt.cm.gray, interpolation='nearest')
        ax0.set_title("Original")
        ax1.imshow(-distance, cmap=plt.cm.jet, interpolation='nearest')
        ax1.set_title("Distance")
        ax2.imshow(markers, cmap=plt.cm.nipy_spectral, interpolation='nearest')
        ax2.set_title("Markers")
        ax3.imshow(labels, cmap=plt.cm.nipy_spectral, interpolation='nearest')
        ax3.set_title("Segmented")

        for ax in axes:
            ax.axis('off')

        fig.tight_layout()
        plt.show()
    def fenShuiLin2(self):
        from scipy import ndimage as ndi
        from skimage import morphology, feature
        from skimage.segmentation import watershed

        x, y = np.indices((80, 80))
        x1, y1, x2, y2 = 28, 28, 44, 52
        r1, r2 = 16, 20
        mask_circle1 = (x - x1) ** 2 + (y - y1) ** 2 < r1 ** 2
        mask_circle2 = (x - x2) ** 2 + (y - y2) ** 2 < r2 ** 2
        image = np.logical_or(mask_circle1, mask_circle2)

        # 绘制原始图片
        plt.subplot(1, 3, 1)
        plt.imshow(image, 'gray')
        plt.title('Original Image')

        # 现在我们用分水岭算法分离两个圆
        distance = ndi.distance_transform_edt(image)  # 距离变换

        # 使用这样的方式寻找峰值
        local_maxi = feature.peak_local_max(distance, footprint=np.ones((3, 3)))

        # 对峰值进行标记
        markers, num_markers = ndi.label(local_maxi)

        # 使用分水岭算法进行分割
        labels = watershed(-distance, mask=image)  # 基于距离变换的分水岭算法

        fig, axes = plt.subplots(nrows=2, ncols=2, figsize=(8, 8))
        axes = axes.ravel()
        ax0, ax1, ax2, ax3 = axes

        ax0.imshow(image, cmap=plt.cm.gray, interpolation='nearest')
        ax0.set_title("Original")
        ax1.imshow(-distance, cmap=plt.cm.jet, interpolation='nearest')
        ax1.set_title("Distance")
        ax2.imshow(markers, cmap=plt.cm.nipy_spectral, interpolation='nearest')
        ax2.set_title("Markers")
        ax3.imshow(labels, cmap=plt.cm.nipy_spectral, interpolation='nearest')
        ax3.set_title("Segmented")

        for ax in axes:
            ax.axis('off')

        fig.tight_layout()
        plt.show()



